from statistics import mode

import numpy as np
import cv2
from rknn.api import RKNN
import json
import time
import paho.mqtt.client as mqtt
import threading

from PIL import Image, ImageFont, ImageDraw
from wz_ai.inference import detect_faces
from wz_ai.inference import draw_text
from wz_ai.inference import draw_bounding_box
from wz_ai.inference import apply_offsets
from wz_ai.inference import load_detection_model
from wz_ai.preprocessor import preprocess_input



print("XXXXXXXXXX   唯众AI实训平台  XXXXXXXXXX")
print(".........      平台启动中    ..........")
print(".........       请等待      ...........")


# parameters for loading data and images
detection_model_path = 'face_reg.xml'

def get_labels(dataset_name):
    if dataset_name == 'imdb':
        return {0: 'woman', 1: 'man'}

gender_labels = get_labels('imdb')

INPUT_SIZE = 48

def load_model(modle_path):
	# Create RKNN object
	rknn = RKNN()

	print("XXXXXXXXXXX      WZ_AI     XXXXXXXXXXX")
	print('-->开始加载模型')
	rknn.load_rknn(modle_path)
	print('模型加载成功')

	# init runtime environment
	print('--> 初始化环境中......')
	ret = rknn.init_runtime()
	if ret != 0:
		print('环境初始化失败')
		exit(ret)
	print('环境初始化完成')
	print("XXXXXXXXXXX      WZ_AI     XXXXXXXXXXX")
	return rknn

def model_predict(rknn, gray_face):
    gender_prediction = rknn.inference(inputs=[gray_face], data_type='float32', data_format='nhwc')

    # perf
    #print('--> Begin evaluate model performance')
    #perf_results = rknn.eval_perf(inputs=[gray_face])
    #print(perf_results)
    #print('done')

    gender_probability = np.max(gender_prediction)
    gender_label_arg = np.argmax(gender_prediction)
    gender_text = gender_labels[gender_label_arg]

    return gender_text, gender_probability

HOST = "139.9.246.47"
PORT = 18883


def on_message(client, userdata, message):
    print(message.topic + " " + ":" + str(message.payload))


def on_connect(client, userdata, flags, rc):
    print("Connected with result code " + str(rc))
    client.subscribe('xingbie')


def client_pb(message: str):
    time_now = time.strftime('%Y-%m-%d %H-%M-%S', time.localtime(time.time()))
    payload = {'msg': '%s' % message, 'time': '%s' % time_now}
    client.publish('xingbie', json.dumps(payload, ensure_ascii=False))
    print('publish success!!')
    return True


gender_text = ''

flag = 1
def  tests(threadName):
    global gender_text
    global flag
    while flag:
        time.sleep(3)
        if gender_text != '':
            client_pb(gender_text)
            gender_text = ''


if __name__ == '__main__':
    # getting input model shapes for inference
    gender_target_size = (INPUT_SIZE, INPUT_SIZE)
    # hyper-parameters for bounding boxes shape
    frame_window = 10
    gender_offsets = (30, 60)

    # starting lists for calculating modes
    gender_window = []

    # loading models
    rknn = load_model('xingbie.rknn')
    face_detection = load_detection_model(detection_model_path)

    # starting video streaming
    # cv2.namedWindow('WZ_AI_Platform ( Please press Esc to exit !)')
    cv2.namedWindow('WZ_AI_Platform ( Please press Esc to exit !)', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('WZ_AI_Platform ( Please press Esc to exit !)', 1920, 1080)
    video_capture = cv2.VideoCapture(0)

    client = mqtt.Client(client_id='whwzzc_xb')
    client.connect('139.9.246.47', 18883, 60)
    client.username_pw_set('admin', 'public')
    client.loop_start()

    t1 = threading.Thread(target=tests, args=('thread 1',))
    t1.start()

    while True:
        frame_start = time.time()
        bgr_image = video_capture.read()[1]
        gray_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
        rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
        faces = detect_faces(face_detection, rgb_image)
        print('人脸识别耗时: %f' % ( time.time() - frame_start) + ' s')

        for face_coordinates in faces:
            start =  time.time()
            x1, x2, y1, y2 = apply_offsets(face_coordinates, gender_offsets)
            rgb_face = rgb_image[y1:y2, x1:x2]
            try:
                rgb_face = cv2.resize(rgb_face, (gender_target_size))
            except:
                continue

            rgb_face = preprocess_input(rgb_face, True)
            rgb_face = np.expand_dims(rgb_face, 0)
            rgb_face = np.expand_dims(rgb_face, -1)
            gender_text, gender_probability = model_predict(rknn, rgb_face)
            print('性别识别耗时: %f'%( time.time() - start)+' s')
            gender_window.append(gender_text)

            if len(gender_window) > frame_window:
                gender_window.pop(0)
            try:
                gender_mode = mode(gender_window)
            except:
                continue

            if gender_text == 'woman':
                color = gender_probability * np.asarray((255, 0, 0))
            elif gender_text == 'man':
                color = gender_probability * np.asarray((0, 255, 255))
            else:
               color = gender_probability * np.asarray((0, 255, 0))

            color = color.astype(int)
            color = color.tolist()
            draw_bounding_box(face_coordinates, rgb_image, color)
            draw_text(face_coordinates, rgb_image, gender_mode,
                      color, 0, -20, 1, 1)
            print('gender_mode=', gender_mode)
        fps_text = ("%.2f" % (1 / ( time.time() - frame_start)))
        cv2.putText(rgb_image, 'FPS: ' +fps_text, (10, 20),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5, (0, 255, 0), 1, cv2.LINE_AA)

        bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
        change_img = Image.fromarray(rgb_image)

        # PIL图片上打印汉字
        draw = ImageDraw.Draw(change_img)  # 图片上打印
        font = ImageFont.truetype("wz_ziti.ttf", 30, encoding="utf-8")  # 参数1：字体文件路径，参数2：字体大小
        draw.text((220, 0), "唯众人脸性别识别系统", (255, 255, 0), font=font)  # 参数1：打印坐标，参数2：文本，参数3：字体颜色，参数4：字体
        font = ImageFont.truetype("wz_ziti.ttf", 20, encoding="utf-8")
        draw.text((215, 40), "(根据人脸性别进行预测并问好)", (255, 0, 0), font=font)
        cv2charimg = cv2.cvtColor(np.array(change_img), cv2.COLOR_RGB2BGR)

        cv2.imshow('WZ_AI_Platform ( Please press Esc to exit !)', cv2charimg)

        '''
        if gender_text != '':
            client_pb(gender_text)
            gender_text = ''
        '''
        if cv2.waitKey(1) & 0xFF == 27:
            flag = 0
            break
